package dao

import (
	"encoding/json"
	"gorm.io/gorm"
	"hhos-mdijounery-server/model"
	"time"
)

type MidjourneyDao interface {
	CreateTaskByImagine(idempotent, prompt, imgUrl string) (*model.MidjourneyImageTask, error)
	CreateTaskByChange(idempotent, action string, index int64) (*model.MidjourneyImageTask, error)
	TaskItem(idempotent string) (*model.MidjourneyImageTask, error)
	TaskItemById(taskId int64) (*model.MidjourneyImageTask, error)
	UpdateTaskItemByIdempotent(taskId string, idempotent string, imageUrl string, paramsJson string, updateData map[string]interface{}) error
	InsertTaskLogItem(taskId string, idempotent string, imageUrl string, paramsJson string) error
}

type MidjourneyDAO struct {
	db *gorm.DB
}

func NewMidjourneyDAO(db *gorm.DB) *MidjourneyDAO {
	return &MidjourneyDAO{
		db: db,
	}
}

// CreateTaskByImagine 创建任务
func (m MidjourneyDAO) CreateTaskByImagine(idempotent, prompt, imgUrl string) (*model.MidjourneyImageTask, error) {
	paramsMap := map[string]string{
		"prompt":  prompt,
		"img_url": imgUrl,
	}

	paramsByte, err := json.Marshal(paramsMap)
	if err != nil {
		return nil, err
	}

	midjourneyTask := &model.MidjourneyImageTask{
		Type:         model.TypeImagine,
		TaskStatus:   model.TaskStatusWaitResponse,
		Idempotent:   idempotent,
		TaskProgress: 0,
		Params:       string(paramsByte),
		CreatedAt:    time.Time{},
		UpdatedAt:    time.Time{},
	}

	err = m.db.Create(midjourneyTask).Error
	return midjourneyTask, err
}

// CreateTaskByChange 绘画改变
func (m MidjourneyDAO) CreateTaskByChange(idempotent, action string, index int64) (*model.MidjourneyImageTask, error) {
	paramsMap := map[string]interface{}{
		"action": action,
		"index":  index,
	}

	paramsByte, err := json.Marshal(paramsMap)
	if err != nil {
		return nil, err
	}

	midjourneyTask := &model.MidjourneyImageTask{
		Type:         model.TypeChange,
		TaskStatus:   model.TaskStatusWaitResponse,
		Idempotent:   idempotent,
		TaskProgress: 0,
		Params:       string(paramsByte),
	}

	err = m.db.Create(midjourneyTask).Error
	return midjourneyTask, err
}

// TaskItem 任务数据
func (m MidjourneyDAO) TaskItem(idempotent string) (*model.MidjourneyImageTask, error) {

	var midjourneyTask model.MidjourneyImageTask
	err := m.db.Where("idempotent=?", idempotent).First(&midjourneyTask).Error
	if err != nil {
		return nil, err
	}

	return &midjourneyTask, err
}

// TaskItemById 任务数据
func (m MidjourneyDAO) TaskItemById(taskId int64) (*model.MidjourneyImageTask, error) {

	var midjourneyTask model.MidjourneyImageTask
	err := m.db.Where("id=?", taskId).First(&midjourneyTask).Error
	if err != nil {
		return nil, err
	}

	return &midjourneyTask, err
}

// UpdateTaskItemByIdempotent 更新任务数据
func (m MidjourneyDAO) UpdateTaskItemByIdempotent(taskId string, idempotent string, imageUrl string, paramsJson string, updateData map[string]interface{}) error {
	err := m.db.Transaction(func(tx *gorm.DB) error {
		var midjourneyImageTask model.MidjourneyImageTask
		err := m.db.Model(&midjourneyImageTask).Where("idempotent=?", idempotent).Updates(updateData).Error
		if err != nil {
			return err
		}
		err = m.InsertTaskLogItem(taskId, idempotent, imageUrl, paramsJson)
		if err != nil {
			return err
		}
		return err
	})

	return err
}

// InsertTaskLogItem 添加任务日志数据
func (m MidjourneyDAO) InsertTaskLogItem(taskId string, idempotent string, imageUrl string, paramsJson string) error {
	var taskLog = &model.MidjourneyImageTaskLog{
		TaskID:     taskId,
		Idempotent: idempotent,
		ImageUrl:   imageUrl,
		Params:     paramsJson,
	}

	err := m.db.Create(taskLog).Error
	return err
}
